package com.leetcode;

import java.util.Random;

/**
 * 215. 数组中的第K个最大元素
 */
public class Solution215_5 {

    private int[] nums;

    public int findKthLargest(int[] nums, int k) {
        this.nums = nums;
        return quick_select(0, nums.length - 1, nums.length - k);
    }

    private int quick_select(int left, int right, int k_max) {
        int idx = random_partition(left, right);
        if (idx == k_max) {
            return nums[idx];
        } else if (idx < k_max) {
            return quick_select(idx + 1, right, k_max);
        } else {
            return quick_select(left, idx - 1, k_max);
        }
    }

    private int random_partition(int left, int right) {
        Random random = new Random();
        int pivot = left + random.nextInt(right - left + 1);
        return partition(left, right, pivot);
    }

    // 三路快速排序
    private int partition(int left, int right, int pivot) {
        int pivot_num = nums[pivot];
        // [left, l]区间的元素 小于 pivot_num
        int l = left - 1;
        // [l + 1, i - 1]区间的元素 等于 pivot_num
        // [r, right]区间的元素 大于 pivot_num
        int r = right + 1;
        int i = left;
        while (i < r) {
            if (nums[i] == pivot_num) {
                i++;
            } else if (nums[i] < pivot_num) {
                swap(++l, i++);
            } else {
                swap(--r, i);
            }
        }
        return i - 1;
    }

    private void swap(int idx1, int idx2) {
        int tmp = nums[idx1];
        nums[idx1] = nums[idx2];
        nums[idx2] = tmp;
    }

    public static void main(String[] args) {
        int[] nums = new int[]{3, 2, 1, 5, 6, 4};
//        int[] nums = new int[]{3};
        int k = 1;
        int num = new Solution215_5().findKthLargest(nums, k);
        System.out.println(num);
    }
}
